import struct
from dataclasses import dataclass
from enum import Enum
from typing import Tuple


class MessageType(Enum):
    PING = 0x0001
    PONG = 0x0002
    DATA = 0x0003
    ACK = 0x0004
    ERROR = 0x0005


@dataclass
class Message:
    msg_type: MessageType
    seq_id: int
    data: bytes

    @property
    def length(self) -> int:
        """计算消息总长度"""
        return 10 + len(self.data)  # 4+2+4 + data_length


class MessageProtocol:
    """
    消息协议处理器

    协议格式:
    +--------+--------+--------+--------+
    | Length | Type   | SeqID  | Data   |
    | 4bytes | 2bytes | 4bytes | N bytes|
    +--------+--------+--------+--------+

    - Length: 消息总长度 (包含头部)
    - Type: 消息类型
    - SeqID: 序列号 (用于请求响应匹配)
    - Data: 消息数据
    """

    HEADER_SIZE = 10  # 4 + 2 + 4
    HEADER_FORMAT = '!IHI'  # 网络字节序: uint32, uint16, uint32

    def encode_message(self, msg_type: MessageType, seq_id: int, data: bytes = b'') -> bytes:
        """
        将消息编码为字节流

        Args:
            msg_type: 消息类型
            seq_id: 序列号
            data: 消息数据

        Returns:
            编码后的字节流
        """
        if not isinstance(msg_type, MessageType):
            raise ValueError(f"Invalid message type: {msg_type}")

        if not isinstance(data, bytes):
            data = str(data).encode('utf-8')

        # 计算总长度
        total_length = self.HEADER_SIZE + len(data)

        # 打包消息头
        header = struct.pack(self.HEADER_FORMAT, total_length, msg_type.value, seq_id)

        # 拼接完整消息
        return header + data

    def decode_message(self, data: bytes) -> Tuple[Message, int]:
        """
        从字节流解码消息

        Args:
            data: 字节流数据

        Returns:
            (解码的消息, 消费的字节数)

        Raises:
            ValueError: 数据格式错误
            struct.error: 解包失败
        """
        if len(data) < self.HEADER_SIZE:
            raise ValueError(f"Data too short for header: {len(data)} < {self.HEADER_SIZE}")

        # 解析消息头
        length, msg_type_value, seq_id = struct.unpack(self.HEADER_FORMAT, data[:self.HEADER_SIZE])

        # 验证消息类型
        try:
            msg_type = MessageType(msg_type_value)
        except ValueError:
            raise ValueError(f"Unknown message type: {msg_type_value}")

        # 验证长度
        if length < self.HEADER_SIZE:
            raise ValueError(f"Invalid message length: {length}")

        if len(data) < length:
            raise ValueError(f"Incomplete message: {len(data)} < {length}")

        # 提取数据部分
        data_length = length - self.HEADER_SIZE
        message_data = data[self.HEADER_SIZE:self.HEADER_SIZE + data_length]

        message = Message(msg_type=msg_type, seq_id=seq_id, data=message_data)
        return message, length

    def create_ping(self, seq_id: int) -> bytes:
        """创建PING消息"""
        return self.encode_message(MessageType.PING, seq_id)

    def create_pong(self, seq_id: int) -> bytes:
        """创建PONG消息"""
        return self.encode_message(MessageType.PONG, seq_id)

    def create_data(self, seq_id: int, data: bytes) -> bytes:
        """创建DATA消息"""
        return self.encode_message(MessageType.DATA, seq_id, data)

    def create_ack(self, seq_id: int) -> bytes:
        """创建ACK消息"""
        return self.encode_message(MessageType.ACK, seq_id)

    def create_error(self, seq_id: int, error_msg: str) -> bytes:
        """创建ERROR消息"""
        return self.encode_message(MessageType.ERROR, seq_id, error_msg.encode('utf-8'))


class MessageBuffer:
    """
    消息缓冲区 - 处理TCP流中的消息边界问题
    """

    def __init__(self):
        self.buffer = bytearray()
        self.protocol = MessageProtocol()

    def feed_data(self, data: bytes) -> list[Message]:
        """
        向缓冲区添加数据并尝试解析消息

        Args:
            data: 新接收的数据

        Returns:
            解析出的完整消息列表
        """
        self.buffer.extend(data)
        messages = []

        while len(self.buffer) >= MessageProtocol.HEADER_SIZE:
            try:
                message, consumed = self.protocol.decode_message(bytes(self.buffer))
                messages.append(message)
                # 移除已消费的数据
                del self.buffer[:consumed]
            except ValueError:
                # 数据不完整，等待更多数据
                break

        return messages

    def clear(self):
        """清空缓冲区"""
        self.buffer.clear()


if __name__ == "__main__":
    # 创建协议处理器
    protocol = MessageProtocol()
    "0000000a000100000001"
    # 编码消息
    ping_msg = protocol.create_ping(seq_id=1)
    data_msg = protocol.create_data(seq_id=2, data=b"Hello, World!")

    print(f"PING message: {ping_msg.hex()}  {ping_msg}")
    print(f"DATA message: {data_msg.hex()}  {data_msg}")

    # 解码消息
    decoded_ping, _ = protocol.decode_message(ping_msg)
    decoded_data, _ = protocol.decode_message(data_msg)

    print(f"Decoded PING: {decoded_ping}")
    print(f"Decoded DATA: {decoded_data}")
    print(f"Data content: {decoded_data.data.decode('utf-8')}")

    # 测试消息缓冲区
    buffer = MessageBuffer()

    # 模拟分片接收
    fragment1 = ping_msg[:2]
    fragment2 = ping_msg[2:] + data_msg

    messages1 = buffer.feed_data(fragment1)
    print(f"Fragment 1 messages: {len(messages1)}")

    messages2 = buffer.feed_data(fragment2)
    print(f"Fragment 2 messages: {len(messages2)}")
    print(f"Total messages: {[msg.msg_type.name for msg in messages2]}")